#include "seccomp_filter.h"
#include <cstring>
#include <stdexcept>
#include <spdlog/spdlog.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/mman.h>  // 添加这个头文件以使用 PROT_EXEC

namespace judge {

SeccompFilter::SeccompFilter(const std::string& language,
                           const std::vector<std::string>& allowed_syscalls,
                           const std::map<std::string, std::map<unsigned int, uint64_t>>& restricted_syscalls)
    : language_(language), allowed_syscalls_(allowed_syscalls), restricted_syscalls_(restricted_syscalls) {
    ctx_ = seccomp_init(SCMP_ACT_KILL);
    if (!ctx_) {
        throw std::runtime_error("Failed to initialize seccomp context");
    }
}

SeccompFilter::~SeccompFilter() {
    if (ctx_) {
        seccomp_release(ctx_);
        ctx_ = nullptr;
    }
}

void SeccompFilter::load() {
    // 先设置网络规则（默认全部禁止，因为默认动作是 KILL）
    setupNetworkRules();
    
    // 再添加允许的系统调用
    setupBasicRules();
    setupLanguageSpecificRules();
    setupFileAccessRules();

    // 加载规则
    if (seccomp_load(ctx_) < 0) {
        throw std::runtime_error("Failed to load seccomp rules");
    }
}

void SeccompFilter::setupBasicRules() {
    // 基本系统调用白名单
    const int basic_syscalls[] = {
        SCMP_SYS(read),
        SCMP_SYS(write),
        SCMP_SYS(close),
        SCMP_SYS(fstat),
        SCMP_SYS(mmap),
        SCMP_SYS(mprotect),
        SCMP_SYS(munmap),
        SCMP_SYS(brk),
        SCMP_SYS(access),
        SCMP_SYS(arch_prctl),
        SCMP_SYS(exit_group),
        SCMP_SYS(exit),
        SCMP_SYS(fcntl),
        SCMP_SYS(lseek),
        SCMP_SYS(getdents64),
        SCMP_SYS(stat),
        SCMP_SYS(fstat),
        SCMP_SYS(lstat),
        SCMP_SYS(statfs),
        SCMP_SYS(fstatfs),
        SCMP_SYS(sysinfo),
        SCMP_SYS(getrlimit),
        SCMP_SYS(setrlimit),
        SCMP_SYS(getrusage),
        SCMP_SYS(clock_getres),
        SCMP_SYS(sched_yield),
        SCMP_SYS(sched_getaffinity),
        SCMP_SYS(pipe2),
        SCMP_SYS(select),
        SCMP_SYS(poll),
        SCMP_SYS(epoll_create1),
        SCMP_SYS(epoll_ctl),
        SCMP_SYS(epoll_wait),
        SCMP_SYS(eventfd2),
        SCMP_SYS(timerfd_create),
        SCMP_SYS(timerfd_settime),
        SCMP_SYS(timerfd_gettime),
        SCMP_SYS(memfd_create),
        SCMP_SYS(ftruncate),
        SCMP_SYS(fallocate),
        SCMP_SYS(set_tid_address),
        SCMP_SYS(set_robust_list),
        SCMP_SYS(rseq),
        SCMP_SYS(madvise),
        SCMP_SYS(pkey_mprotect),
        SCMP_SYS(pkey_alloc),
        SCMP_SYS(pkey_free),
        SCMP_SYS(sigaltstack),
        SCMP_SYS(prctl),
        SCMP_SYS(nanosleep),
        SCMP_SYS(clock_nanosleep),
        SCMP_SYS(getpid),
        SCMP_SYS(gettid),
        SCMP_SYS(tgkill),
        SCMP_SYS(futex),
        SCMP_SYS(sched_getaffinity),
        SCMP_SYS(sched_yield),
        SCMP_SYS(rt_sigaction),
        SCMP_SYS(rt_sigprocmask),
        SCMP_SYS(rt_sigreturn),
        SCMP_SYS(restart_syscall)
    };

    for (int syscall : basic_syscalls) {
        int rc = seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, syscall, 0);
        if (rc < 0) {
            spdlog::error("Failed to add basic syscall rule for syscall {}: {}", syscall, strerror(-rc));
            throw std::runtime_error("Failed to add basic syscall rule");
        }
    }

    // 限制mmap的权限
    if (seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, SCMP_SYS(mmap), 1,
                        SCMP_CMP(2, SCMP_CMP_MASKED_EQ, PROT_EXEC, 0)) < 0) {
        throw std::runtime_error("Failed to add mmap rule");
    }

    // 限制mprotect的权限
    if (seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, SCMP_SYS(mprotect), 1,
                        SCMP_CMP(2, SCMP_CMP_MASKED_EQ, PROT_EXEC, 0)) < 0) {
        throw std::runtime_error("Failed to add mprotect rule");
    }

    // 明确禁止一些危险的系统调用
    const int dangerous_syscalls[] = {
        SCMP_SYS(fork),
        SCMP_SYS(vfork),
        SCMP_SYS(clone),
        SCMP_SYS(clone3),
        SCMP_SYS(execve),
        SCMP_SYS(execveat),
        SCMP_SYS(ptrace),
        SCMP_SYS(process_vm_readv),
        SCMP_SYS(process_vm_writev),
        SCMP_SYS(kill),
        SCMP_SYS(tkill),
        SCMP_SYS(tgkill),
        SCMP_SYS(mount),
        SCMP_SYS(umount2),
        SCMP_SYS(pivot_root),
        SCMP_SYS(chroot),
        SCMP_SYS(setns),
        SCMP_SYS(unshare)
    };

    for (int syscall : dangerous_syscalls) {
        int rc = seccomp_rule_add(ctx_, SCMP_ACT_KILL, syscall, 0);
        if (rc < 0) {
            spdlog::error("Failed to add kill rule for dangerous syscall {}: {}", syscall, strerror(-rc));
            throw std::runtime_error("Failed to add dangerous syscall rule");
        }
    }

    // 禁止所有 shell 相关的程序
    const char* shell_programs[] = {
        "/bin/sh",
        "/bin/bash",
        "/bin/dash",
        "/bin/rbash",
        "/bin/zsh",
        "/usr/bin/sh",
        "/usr/bin/bash",
        "/usr/bin/dash",
        "/usr/bin/rbash",
        "/usr/bin/zsh"
    };

    for (const char* shell : shell_programs) {
        if (seccomp_rule_add(ctx_, SCMP_ACT_KILL, SCMP_SYS(execve), 1,
                            SCMP_CMP(0, SCMP_CMP_EQ, (scmp_datum_t)shell)) < 0) {
            spdlog::error("Failed to add kill rule for shell: {}", shell);
            throw std::runtime_error("Failed to add shell kill rule");
        }
    }

    // 禁止所有系统管理命令
    const char* system_commands[] = {
        "/bin/ps",
        "/usr/bin/ps",
        "/bin/top",
        "/usr/bin/top",
        "/bin/netstat",
        "/usr/bin/netstat",
        "/bin/ifconfig",
        "/usr/sbin/ifconfig",
        "/bin/ip",
        "/usr/bin/ip"
    };

    for (const char* cmd : system_commands) {
        if (seccomp_rule_add(ctx_, SCMP_ACT_KILL, SCMP_SYS(execve), 1,
                            SCMP_CMP(0, SCMP_CMP_EQ, (scmp_datum_t)cmd)) < 0) {
            spdlog::error("Failed to add kill rule for command: {}", cmd);
            throw std::runtime_error("Failed to add command kill rule");
        }
    }
}

void SeccompFilter::setupLanguageSpecificRules() {
    if (language_ == "cpp" || language_ == "c") {
        // C/C++特定的系统调用
        const int cpp_syscalls[] = {
            SCMP_SYS(madvise),
            SCMP_SYS(prlimit64),
            SCMP_SYS(set_mempolicy),
            SCMP_SYS(mbind),
            SCMP_SYS(migrate_pages),
            SCMP_SYS(move_pages)
        };
        for (int syscall : cpp_syscalls) {
            if (seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, syscall, 0) < 0) {
                throw std::runtime_error("Failed to add C/C++ syscall rule");
            }
        }
    } else if (language_ == "python") {
        // Python特定的系统调用
        const int python_syscalls[] = {
            SCMP_SYS(getdents64),
            SCMP_SYS(lseek),
            SCMP_SYS(ioctl),
            SCMP_SYS(fcntl),
            SCMP_SYS(stat),
            SCMP_SYS(getcwd),
            SCMP_SYS(sysinfo),
            SCMP_SYS(getrlimit),
            SCMP_SYS(prlimit64),
            SCMP_SYS(sched_getaffinity),
            SCMP_SYS(getrandom)
        };
        for (int syscall : python_syscalls) {
            if (seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, syscall, 0) < 0) {
                throw std::runtime_error("Failed to add Python syscall rule");
            }
        }
    }

    // 添加用户指定的系统调用
    for (const auto& syscall : allowed_syscalls_) {
        int syscall_nr = seccomp_syscall_resolve_name(syscall.c_str());
        if (syscall_nr == __NR_SCMP_ERROR) {
            spdlog::warn("Unknown syscall: {}", syscall);
            continue;
        }
        if (seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, syscall_nr, 0) < 0) {
            throw std::runtime_error("Failed to add user-specified syscall rule");
        }
    }

    // 添加带限制的系统调用
    for (const auto& [syscall, args] : restricted_syscalls_) {
        int syscall_nr = seccomp_syscall_resolve_name(syscall.c_str());
        if (syscall_nr == __NR_SCMP_ERROR) {
            spdlog::warn("Unknown syscall: {}", syscall);
            continue;
        }

        std::vector<struct scmp_arg_cmp> arg_cmps;
        for (const auto& [arg_num, value] : args) {
            arg_cmps.push_back(SCMP_CMP(arg_num, SCMP_CMP_EQ, value));
        }

        if (seccomp_rule_add_array(ctx_, SCMP_ACT_ALLOW, syscall_nr,
                                 arg_cmps.size(), arg_cmps.data()) < 0) {
            throw std::runtime_error("Failed to add restricted syscall rule");
        }
    }
}

void SeccompFilter::setupFileAccessRules() {
    // 文件访问相关的系统调用，带参数限制
    const int file_syscalls[] = {
        SCMP_SYS(openat),
        SCMP_SYS(open)
    };

    for (int syscall : file_syscalls) {
        if (seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, syscall, 0) < 0) {
            throw std::runtime_error("Failed to add open/openat rule");
        }
    }

    // 允许删除临时文件
    if (seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, SCMP_SYS(unlink), 0) < 0) {
        throw std::runtime_error("Failed to add file deletion rule");
    }
    if (seccomp_rule_add(ctx_, SCMP_ACT_ALLOW, SCMP_SYS(unlinkat), 0) < 0) {
        throw std::runtime_error("Failed to add file deletion rule");
    }
}

void SeccompFilter::setupNetworkRules() {
    // 由于默认动作是 KILL，所以这里不需要显式添加 KILL 规则
    // 网络相关的系统调用会被默认 KILL
    spdlog::info("Network syscalls will be killed by default");
}

} // namespace judge 